import numpy
import numpy as np
import cv2
from PyVGrid.vgrid import BodyPart

def exposure_fusion(sequence, refer_illumination_id=3, best_illumination=0.5, sigma=0.2, layers_num=5, boday_part=BodyPart().cardiac, tqdm_bar=None):
    # 转化成 float 数据
    sequence = numpy.stack([it.astype("float32") / 65535 for it in sequence], axis=0)
    S = len(sequence)
    H, W = sequence[0].shape
    # 准备一些中间变量
    laplace_kernel = numpy.array(([0, 1, 0], [1, -4, 1], [0, 1, 0]), dtype="float32")
    mse = lambda l, r: (l - r) * (l - r)
    best_illumination = numpy.full((H, W), best_illumination, dtype='float32')

    if boday_part == BodyPart().cardiac:
        print("Warning: Body part is set to cardiac")
        pass
    elif boday_part == BodyPart().neurolgy:
        best_illumination[sequence[refer_illumination_id, :, :] < 1] = 0.75                  # TODO # 神经参数   sub参数2  b6作为基准
        best_illumination[sequence[refer_illumination_id, :, :] < 0.5] = 0.5                 # TODO
        best_illumination[sequence[refer_illumination_id, :, :] < 0.25] = 0.25               # TODO
        contrasts_weight, illuminations_weight = 2, 1
    elif boday_part == BodyPart().Peripheral:
        best_illumination[sequence[refer_illumination_id, :, :] < 1] = 0.65                  # TODO # 下肢参数   sub参数3   b6作为基准
        best_illumination[sequence[refer_illumination_id, :, :] < 0.5] = 0.5                 # TODO
        best_illumination[sequence[refer_illumination_id, :, :] < 0.25] = 0.25               # TODO
        contrasts_weight, illuminations_weight = 1, 1
    elif boday_part == BodyPart().thorax:
        best_illumination[sequence[refer_illumination_id, :, :] < 1] = 0.65                  # TODO  # 胸部和腹部通用参数   sub参数3  b6还是b5有待商榷
        best_illumination[sequence[refer_illumination_id, :, :] < 0.5] = 0.3                 # TODO
        best_illumination[sequence[refer_illumination_id, :, :] < 0.25] = 0.2               # TODO
        contrasts_weight, illuminations_weight = 2, 1
    elif boday_part == BodyPart().abdomen:
        best_illumination[sequence[refer_illumination_id, :, :] < 1] = 0.65                  # TODO  # 胸部和腹部通用参数   sub参数3  b6还是b5有待商榷
        best_illumination[sequence[refer_illumination_id, :, :] < 0.5] = 0.3                 # TODO
        best_illumination[sequence[refer_illumination_id, :, :] < 0.25] = 0.2               # TODO
        contrasts_weight, illuminations_weight = 2, 1
    else:
        print("Warning: unknown bodypart")


    # 存放每张图像的权重图
    normalize = lambda x: x / numpy.expand_dims(numpy.sum(x, axis=0), axis=0)
    contrasts, illuminations = [], []
    for s in range(S):
        # 从拉普拉斯求对比度
        contrast = cv2.filter2D(sequence[s], -1, laplace_kernel, borderType=cv2.BORDER_REPLICATE)
        contrast = numpy.abs(contrast)
        contrasts.append(contrast)
        # 求亮度
        illumination = [numpy.exp(-0.5 * mse(sequence[s][:, :], best_illumination) / (sigma * sigma))]
        illumination = numpy.prod(illumination, axis=0)
        illuminations.append(illumination)
    contrasts = np.array(contrasts)
    contrasts = contrasts / numpy.sum(contrasts, axis=0)
    contrasts = contrasts * contrasts_weight

    illuminations = np.array(illuminations)
    illuminations = illuminations / numpy.sum(illuminations, axis=0)
    illuminations = illuminations * illuminations_weight

    weights = contrasts + illuminations
    weights = normalize(weights)
    # print(contrasts[:, 26, 553], illuminations[:, 26, 553], weights[:, 26, 553])
    # 这里要把 sequence 还原回来
    sequence *= 65535
    # 根据最高分辨率的图像 high_res, 得到高度 layers 的高斯金字塔
    def build_gaussi_pyramid(high_res, layers):
        this_flash = [high_res]
        for i in range(1, layers):
            # 先对当前权重做高斯模糊, 然后下采样
            blurred = cv2.GaussianBlur(this_flash[i - 1], (5, 5), 0.83)
            blurred = blurred[::2, ::2]
            this_flash.append(blurred)
        return this_flash

    # 根据已知的高斯金字塔, 从最底层开始上采样, 得到每一个尺度的 laplace 细节
    def build_laplace_pyramaid(gaussi_pyramid, layers):
        upsampled = gaussi_pyramid[layers - 1]
        pyramid = [upsampled]
        for i in range(layers - 1, 0, -1):
            size = (gaussi_pyramid[i - 1].shape[1], gaussi_pyramid[i - 1].shape[0])
            upsampled = cv2.resize(gaussi_pyramid[i], size)
            pyramid.append(gaussi_pyramid[i - 1] - upsampled)
        pyramid.reverse()  # 目前分辨率都是从高到低排列的
        return pyramid

    # 求每张图的权重的高斯金字塔
    sequence_weights_pyramids = [build_gaussi_pyramid(weights[s], layers_num) for s in range(S)]
    # 求每张图的高斯金字塔, 以求 laplace
    sequence_gaussi_pyramids = [build_gaussi_pyramid(sequence[s], layers_num) for s in range(S)]
    # 求每张图的 laplace 金字塔
    sequence_laplace_pyramids = [build_laplace_pyramaid(sequence_gaussi_pyramids[s], layers_num) for s in range(S)]
    # 这里可以归一化
    # 每一个尺度, 融合一系列图像的的 laplace 细节, 得到一个融合的 laplace 金字塔
    fused_laplace_pyramid = [numpy.sum([sequence_laplace_pyramids[k][n] * sequence_weights_pyramids[k][n] for k in range(S)], axis=0) for n in range(layers_num)]
    # fused_laplace_pyramid = [numpy.sum([sequence_laplace_pyramids[k][n] *
    #                                     numpy.atleast_3d(sequence_weights_pyramids[k][n]) for k in range(S)], axis=0) for n in range(layers_num)]

    # 先从最底层的图像开始, 每次上采样都加上同等尺度的 laplace 细节
    start = fused_laplace_pyramid[layers_num - 1]
    for i in range(layers_num - 2, -1, -1):
        upsampled = cv2.resize(start, (fused_laplace_pyramid[i].shape[1], fused_laplace_pyramid[i].shape[0]))
        start = fused_laplace_pyramid[i] + upsampled
    # 灰度值截断在 0-255 之间
    start = numpy.clip(start, 0, 65535).astype("uint16")
        # 放到结果列表中

    return start